#!/usr/bin/env python
# -*- coding: utf-8 -*-

import asyncio
import uvicorn
import argparse
import logging
import sys
import pymysql
from pymysql.err import OperationalError

from app.core.config import settings
from app.core.db_init import initialize_database
from app.repositories.milvus import milvus_client
from app.repositories.vector import VectorRepository


def create_database():
    """创建数据库"""
    try:
        # 连接到MySQL服务器
        conn = pymysql.connect(
            host=settings.MYSQL_HOST,
            port=settings.MYSQL_PORT,
            user=settings.MYSQL_USER,
            password=settings.MYSQL_PASSWORD
        )
        
        # 创建游标
        cursor = conn.cursor()
        
        # 创建数据库
        cursor.execute(f"CREATE DATABASE IF NOT EXISTS {settings.MYSQL_DB}")
        
        # 提交事务
        conn.commit()
        
        # 关闭连接
        cursor.close()
        conn.close()
        
        print(f"数据库 {settings.MYSQL_DB} 创建成功")
    except Exception as e:
        print(f"创建数据库失败: {str(e)}")
        raise


async def setup():
    """初始化应用程序"""
    # 创建数据库
    create_database()
    
    # 初始化数据库
    await initialize_database()
    
    # 初始化Milvus
    try:
        await milvus_client.connect()
        
        # 断开Milvus连接
        await milvus_client.disconnect()
    except Exception as e:
        print(f"Milvus初始化失败: {str(e)}")
        print("继续初始化过程...")


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="运行Text2SQL RAG API")
    parser.add_argument("--host", default="0.0.0.0", help="主机地址")
    parser.add_argument("--port", type=int, default=8000, help="端口号")
    parser.add_argument("--reload", action="store_true", help="是否启用热重载")
    parser.add_argument("--setup", action="store_true", help="是否初始化应用程序")
    parser.add_argument("--init-only", action="store_true", help="仅初始化，不启动服务器")
    args = parser.parse_args()
    
    # 初始化应用程序
    if args.setup:
        asyncio.run(setup())
        print("应用程序初始化完成")
        
        if args.init_only:
            return
    
    print(f"正在启动服务器，地址: {args.host}:{args.port}")
    
    # 运行应用程序
    uvicorn.run(
        "app.main:app",
        host=args.host,
        port=args.port,
        reload=args.reload,
        log_level="info"
    )


if __name__ == "__main__":
    main()